import os
import json
import requests
from openai import OpenAI
from dotenv import load_dotenv, find_dotenv

_ = load_dotenv(find_dotenv())
amap_key = os.environ.get('amap_key')
client = OpenAI()

# 封装地方天气查询函数
def get_wx(adcode):
    url = f"""https://restapi.amap.com/v3/weather/weatherInfo?key={
        amap_key}&city={adcode}"""
    ret = requests.get(url)
    result = ret.json() 
    if "lives" in result and result["lives"]:
        return result["lives"]
    return None

# 创建一个用于Function Calling的函数定义
def create_function_calling_request(prompt):
    return client.chat.completions.create(
        model="deepseek-chat",
        messages=[
            {"role": "user", "content": prompt}
        ],
      # 定义可用的工具函数列表，供模型在需要时调用
        tools=[
            {   
                "type": "function",  # 工具类型为函数
                "function": {
                    "name": "get_wx",  # 函数名称，必须与实际函数名一致
                    "description": "取得指定区域的天气",  # 函数功能描述，帮助模型理解何时调用
                    "parameters": {         # 函数参数定义
                        "type": "object",   # 参数类型为对象
                        "properties": {     # 参数属性定义
                            "adcode": {     # 参数名
                                "type": "string",  # 参数类型
                                "description": "希望查询天气的区域名称",  # 参数描述
                            }
                        },
                        "required": ["adcode"],  # 必需参数列表
                    }
                }
            }
        ],
        function_call="auto"    # 允许模型自主决定是否调用函数
    )

# 创建一个包含函数调用的prompt
prompt = "北京的天气如何？"

# 发送请求并获取响应
response = create_function_calling_request(prompt)
#print(f"===Debug===\n{response}\n")

# 解析响应并调用相应的函数
function_call = response.choices[0].message.tool_calls
function_name = function_call[0].function.name
function_args = json.loads(function_call[0].function.arguments)

if function_name == "get_wx":
    adcode = function_args.get("adcode")
    result = get_wx(adcode)
    
    # 将本地计算结果返回大模型
    follow_up_prompt = {
        "role": "tool",
        "name": "get_wx",
        "tool_call_id": function_call[0].id,
        "content": json.dumps({"result": result})
    }
    
    # 使用tool_call.id向大模型返回结果
    final_response = client.chat.completions.create(
        model="deepseek-chat",
        messages=[
            {"role": "user", "content": prompt},
            response.choices[0].message,
            follow_up_prompt
        ]
    )
    #print(f"===Debug===\n{final_response}\n")
    # 输出最终结果
    print(final_response.choices[0].message.content.strip())
else:
    print("Function not recognized.")
